package com.wei.czz.framework.common.service;

import com.wei.czz.common.constant.Constant;
import com.wei.czz.common.utils.ContextUtils;
import lombok.AllArgsConstructor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.*;
import java.util.function.Function;
import java.util.function.Supplier;

/**
 * Created by IntelliJ IDEA.
 *
 * @author wei
 * date: 2023-03-31 18:01:25
 * className: AsyncService 异步操作服务接口
 * version: 1.0
 * description:
 */
@Component
@AllArgsConstructor
public class AsyncService {

    private static final Logger log = LoggerFactory.getLogger(AsyncService.class);

    private final ScheduledThreadPoolExecutor delayExecutor;

    private final ThreadPoolTaskExecutor threadExecutor;

    private final SnowflakeService snowflakeService;

    /**
     * 异步执行任务
     * @param runnable 任务对象
     */
    public void execute(Runnable runnable) {
        threadExecutor.execute(runnableDecorator(runnable, true));
    }

    /**
     * 异步执行任务
     * @param runnable 任务对象
     * @return 没有执行结果的CompletableFuture对象
     */
    public CompletableFuture<Void> runAsync(Runnable runnable) {
        return CompletableFuture.runAsync(runnableDecorator(runnable, false), threadExecutor)
                .exceptionally(this.throwableDecorator(e -> {
                    log.error("异步执行异常。message={}", e.getMessage(), e);
                    return null;
                }));
    }

    /**
     * 异步执行任务
     * @param supplier 任务对象
     * @return 执行结果
     * @param <V> 泛型参数
     */
    public <V> CompletableFuture<V> supplyAsync(Supplier<V> supplier) {
        CompletableFuture<V> future = CompletableFuture.supplyAsync(supplierDecorator(supplier), threadExecutor);
        future.exceptionally(this.throwableDecorator(e -> {
            log.error("异步执行异常：message={}", e.getMessage(), e);
            return null;
        }));
        return future;
    }

    /**
     * 异步执行任务
     * @param callable 任务对象
     * @return 执行结果
     * @param <V> 泛型参数
     */
    public <V> Future<V> submit(Callable<V> callable) {
        return threadExecutor.submit(callableDecorator(callable));
    }

    /**
     * 异步延迟执行任务
     * @param runnable 任务对象
     * @param delay    延迟时长
     * @param unit     时长单位
     */
    public void schedule(Runnable runnable, long delay, TimeUnit unit) {
        delayExecutor.schedule(runnableDecorator(runnable, true), delay, unit);
    }

    /**
     * 异步延迟执行任务
     * @param callable 任务对象
     * @param delay    延迟时长
     * @param unit     时长单位
     * @return 执行结果
     * @param <V> 泛型参数
     */
    public <V> ScheduledFuture<V> schedule(Callable<V> callable, long delay, TimeUnit unit) {
        return delayExecutor.schedule(callableDecorator(callable), delay, unit);
    }

    /**
     * Runnable装饰
     * @param runnable 目标Runnable对象
     * @return 装饰后的Runnable对象
     */
    private Runnable runnableDecorator(Runnable runnable, boolean showErrorLog) {
        // 获取父线程的链路追踪配置
        Map<String, String> map = MDC.getCopyOfContextMap();
        map.put(Constant.SPAN_ID, Long.toHexString(snowflakeService.get()));
        // 获取父线程的上下文参数
        Map<String, Object> contextMap = ContextUtils.getCopy();

        return () -> {
            try {
                // 设置子线程链路追踪配置
                MDC.setContextMap(map);
                // 设置子线程的上下文参数
                ContextUtils.set(contextMap);
                // 运行目标线程
                runnable.run();
            } catch (Exception e) {
                if (showErrorLog) {
                    log.error("异步执行异常。message={}", e.getMessage(), e);
                    return;
                }
                throw e;
            } finally {
                // 清理子线程的链路追踪配置
                MDC.clear();
                // 清理子线程的上下文参数
                ContextUtils.remove();
            }
        };
    }

    /**
     * Callable装饰
     * @param callable 目标Callable对象
     * @return 装饰后的Callable对象
     */
    private <V> Callable<V> callableDecorator(Callable<V> callable) {
        // 获取父线程的链路追踪配置
        Map<String, String> map = MDC.getCopyOfContextMap();
        map.put(Constant.SPAN_ID, Long.toHexString(snowflakeService.get()));
        // 获取父线程的上下文参数
        Map<String, Object> contextMap = ContextUtils.getCopy();

        return () -> {
            try {
                // 设置子线程链路追踪配置
                MDC.setContextMap(map);
                // 设置子线程的上下文参数
                ContextUtils.set(contextMap);
                // 运行目标线程
                return callable.call();
            } finally {
                // 清理子线程的链路追踪配置
                MDC.clear();
                // 清理子线程的上下文参数
                ContextUtils.remove();
            }
        };
    }

    /**
     * Supplier装饰
     * @param supplier 目标Supplier对象
     * @return 装饰后的Supplier对象
     */
    private <V> Supplier<V> supplierDecorator(Supplier<V> supplier) {
        // 获取父线程的链路追踪配置
        Map<String, String> map = MDC.getCopyOfContextMap();
        map.put(Constant.SPAN_ID, Long.toHexString(snowflakeService.get()));
        // 获取父线程的上下文参数
        Map<String, Object> contextMap = ContextUtils.getCopy();
        return () -> {
            try {
                // 设置子线程链路追踪配置
                MDC.setContextMap(map);
                // 设置子线程的上下文参数
                ContextUtils.set(contextMap);
                // 运行目标线程
                return supplier.get();
            } finally {
                // 清理子线程的链路追踪配置
                MDC.clear();
                // 清理子线程的上下文参数
                ContextUtils.remove();
            }
        };
    }

    /**
     * throwable装饰
     * @param function 目标Function对象
     * @return 装饰后的Function对象
     * @param <V> 泛型参数
     */
    private <V> Function<Throwable, V> throwableDecorator(Function<Throwable, V> function) {
        // 获取父线程的链路追踪配置
        Map<String, String> map = MDC.getCopyOfContextMap();
        // 获取父线程的上下文参数
        Map<String, Object> contextMap = ContextUtils.getCopy();

        return throwable -> {
            map.put(Constant.SPAN_ID, Long.toHexString(snowflakeService.get()));
            try {
                // 设置子线程链路追踪配置
                MDC.setContextMap(map);
                // 设置子线程的上下文参数
                ContextUtils.set(contextMap);

                return function.apply(throwable);

            } finally {
                // 清理子线程的链路追踪配置
                MDC.clear();
                // 清理子线程的上下文参数
                ContextUtils.remove();
            }
        };
    }

}
